from langchain.chains import LLMChain
from langchain.prompts.chat import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
from langchain_openai import ChatOpenAI

from ai_translator_v2.utils import LOG


class TranslationChain:
    """
    TranslationChain 类用于创建一个语言翻译链，该链结合了OpenAI的聊天模型和LangChain框架，
    用于实现从一种语言到另一种语言的翻译任务。

    参数:
    - model_name: str, 指定使用的OpenAI模型名称，默认为"gpt-3.5-turbo"。
    - verbose: bool, 是否在执行过程中输出详细信息，默认为True。
    """

    def __init__(self, style: str = '翻译家', model_name: str = "gpt-3.5-turbo", verbose: bool = True):
        # 定义系统消息模板，用于引导翻译流程
        template = (
            """You are a translation expert, proficient in various languages, and have a very strong style of translation. 
             Now translates {source_language} to {target_language}, 
             Please translation use {role} style,
             Don't change punctuation marks!
            """
        )
        system_message_prompt = SystemMessagePromptTemplate.from_template(template)

        # 定义人类消息模板，用于接收待翻译文本
        human_template = "{text}"
        human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)

        # 将系统消息和人类消息模板组合成一个聊天提示模板
        chat_prompt_template = ChatPromptTemplate.from_messages([
            system_message_prompt,
            human_message_prompt,
        ])

        # 初始化基于指定模型名称和温度的ChatOpenAI对象
        chat = ChatOpenAI(model_name=model_name, temperature=0, verbose=verbose)

        # 创建LLMChain对象，配置聊天模型和提示模板
        self.chain = LLMChain(llm=chat, prompt=chat_prompt_template, verbose=verbose)

    def run(self, text: str, role: str, source_language: str, target_language: str) -> (str, bool):
        """
        执行文本翻译任务。

        参数:
        - text: str, 待翻译的文本。
        - source_language: str, 源语言代码。
        - target_language: str, 目标语言代码。

        返回:
        - result: str, 翻译后的文本。
        - success: bool, 翻译任务是否成功执行。
        """
        result = ""
        try:
            # 调用chain的run方法执行翻译，并传入必要参数
            result = self.chain.run({
                "text": text,
                "role": role,
                "source_language": source_language,
                "target_language": target_language
            })
        except Exception as e:
            # 如果在翻译过程中出现异常，记录错误信息
            LOG.error(f"An error occurred during translation: {e}")
            return result, False
        return result, True
